# NOTE: Building this image requires docker version >= 18.0

ARG BASE_IMAGE=ubuntu:22.04

FROM ${BASE_IMAGE} AS official

ARG TARGETPLATFORM=linux/arm64

ENV USER_PASSWD=change_me

SHELL [ "/bin/bash", "-c" ]

RUN cp /etc/apt/sources.list /etc/apt/sources.list.backup && \
    case ${TARGETPLATFORM} in \
         "linux/arm64")  sed -i 's|ports.ubuntu.com|mirrors.tuna.tsinghua.edu.cn|g' /etc/apt/sources.list  ;; \
         *)              sed -i 's|archive.ubuntu.com|mirrors.tuna.tsinghua.edu.cn|g' /etc/apt/sources.list  ;; \
    esac

RUN apt-get update \
    && apt-get install --no-install-recommends -y \
        apt-transport-https \
        ca-certificates \
        build-essential \
        bash \
        curl \
        git \
        wget \
        gcc \
        g++ \
        make \
        cmake \
        python3 \
        python3-pip \
        gdb \
        vim \
        file \
        man \
        sudo \
        zlib1g \
        openssl \
        unzip \
        pciutils \
        net-tools \
        gfortran \
        patchelf \
        libblas3 \
        libblas-dev \
        libssl-dev \
        zlib1g-dev \
        libncurses5-dev \
        libbz2-dev \
        libreadline-dev \
        libsqlite3-dev \
        libffi-dev \
        libnss3-dev \
        libgdbm-dev \
        liblzma-dev \
        libev-dev \
        openssh-server \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/* \
    && rm -rf /var/tmp/* \
    && rm -rf /tmp/*

# 创建 hccl 用户
RUN groupadd -g 1000 hcclgroup && \
    useradd -u 1000 -g hcclgroup -ms /bin/bash hccluser && \
    usermod -aG sudo hccluser && \
    echo "hccluser ALL=(ALL) NOPASSWD:/usr/bin/apt-get,/usr/bin/apt" >> /etc/sudoers

USER hccluser 

RUN pip install --no-cache-dir -i https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple \
        attrs cython numpy==1.24.0 decorator sympy cffi pyyaml pathlib2 \
        psutil protobuf==3.20 scipy requests absl-py

# 安装 CANN 8.2.RC1.alpha003
RUN case ${TARGETPLATFORM} in \
         "linux/arm64")  ARCH=aarch64  ;; \
         *)              ARCH=x86_64   ;; \
    esac && \
    CANN_TOOLKIT_URL="https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/Milan-ASL/Milan-ASL%20V100R001C22B800TP052/Ascend-cann-toolkit_8.2.RC1.alpha003_linux-${ARCH}.run" && \
    curl -fsSL -o /tmp/Ascend-cann-toolkit.run -O "${CANN_TOOLKIT_URL}" && \
    CANN_COMMUNITY_SDK_URL="https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/Milan-ASL/Milan-ASL%20V100R001C22B800TP052/Ascend-cann-communitysdk_8.2.RC1.alpha003_linux-${ARCH}.run" && \
    curl -fsSL -o /tmp/Ascend-cann-communitysdk.run -O "${CANN_COMMUNITY_SDK_URL}"

RUN chmod +x /tmp/Ascend-cann-toolkit.run && \
    /tmp/Ascend-cann-toolkit.run --quiet --install && \
    rm /tmp/Ascend-cann-toolkit.run

RUN chmod +x /tmp/Ascend-cann-communitysdk.run && \
    /tmp/Ascend-cann-communitysdk.run --quiet --full && \
    rm /tmp/Ascend-cann-communitysdk.run

# 安装 HCCL 依赖
RUN curl -fsSL -o /tmp/include.zip -O https://github.com/nlohmann/json/releases/download/v3.11.2/include.zip && \
    unzip -d ${HOME}/nlohmann_json /tmp/include.zip && \
    rm /tmp/include.zip

RUN curl -fsSL -o /tmp/mpich.tar.gz -O https://www.mpich.org/static/downloads/3.2.1/mpich-3.2.1.tar.gz && \
    tar -zxf /tmp/mpich.tar.gz -C /tmp && \
    cd /tmp/mpich-3.2.1 && \
    ./configure --disable-fortran --prefix=${HOME}/mpich --with-device=ch3:nemesis && \
    make && make install && \
    rm -r /tmp/mpich-3.2.1 && \
    rm /tmp/mpich.tar.gz

# 设置环境变量
RUN \
    # NPU 驱动环境变量
    echo 'export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/driver:${LD_LIBRARY_PATH}' >> ${HOME}/.bashrc && \
    echo 'export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/common:${LD_LIBRARY_PATH}' >> ${HOME}/.bashrc && \
    # CANN Toolkit 环境变量
    echo 'source ${HOME}/Ascend/ascend-toolkit/set_env.sh' >> ${HOME}/.bashrc && \
    # MPICH 环境变量
    echo 'export PATH=${HOME}/mpich/bin:${PATH}' >> ${HOME}/.bashrc && \
    echo 'export LD_LIBRARY_PATH=${HOME}/mpich/lib:${LD_LIBRARY_PATH}' >> ${HOME}/.bashrc

USER root

# SSH 配置
RUN echo "PasswordAuthentication yes" >> /etc/ssh/sshd_config && \
    echo "PermitRootLogin no" >> /etc/ssh/sshd_config && \
    echo "PermitUserEnvironment yes" >> /etc/ssh/sshd_config && \
    echo "ClientAliveInterval 60" >> /etc/ssh/sshd_config && \
    echo "ClientAliveCountMax 3" >> /etc/ssh/sshd_config && \
    echo "AllowUsers hccluser" >> /etc/ssh/sshd_config

# SSH 启动脚本
RUN echo '#!/bin/bash' > /start.sh && \
    echo 'if [ -n "${USER_PASSWD}" ]; then' >> /start.sh && \
    echo '  echo "hccluser:${USER_PASSWD}" | chpasswd' >> /start.sh && \
    echo 'fi' >> /start.sh && \
    echo 'mkdir -p /var/run/sshd' >> /start.sh && \
    echo 'ssh-keygen -A' >> /start.sh && \
    echo '/usr/sbin/sshd -D -e' >> /start.sh && \
    chmod +x /start.sh

EXPOSE 22

CMD [ "/start.sh" ]
